# Create your views here.
from django.http import JsonResponse
from pymongo import MongoClient
import re
import json
from transformers import AutoTokenizer, AutoModel

tokenizer = AutoTokenizer.from_pretrained("E:\\glm_0429\\ChatGLM-6B-main\\THUDM\\chatglm-6b",
                                          trust_remote_code=True, revision="")
model = AutoModel.from_pretrained("E:\\glm_0429\\ChatGLM-6B-main\\THUDM\\chatglm-6b", trust_remote_code=True,
                                  revision="").half().cuda()




def login(request):
    data = request.POST
    username = data['username']
    password = data['password']
    role = data['role']
    client = MongoClient('mongodb://localhost:27017/')
    db = client['medic_app']
    if role == '0':
        collection = db['doctor']
        if collection.find_one({'id': username}):
            user = collection.find_one({'id': username})
            if user.get('pwd') != password:
                return JsonResponse({"permission": "False"})
            else:
                return JsonResponse({"permission": "True"})
        else:
            return JsonResponse({"permission":"Null"})
    elif role == '1':
        collection = db['patient']
        if collection.find_one({'id': username}):
            user = collection.find_one({'id': username})
            if user.get('pwd') != password:
                return JsonResponse({"permission": "False"})
            else:
                return JsonResponse({"permission": "True"})
        else:
            return JsonResponse({"permission":"Null"})
    else:
        return JsonResponse({"permission": "False"})


def register(request):
    data = request.POST
    username = data['username']
    password = data['password']
    role = data['role']
    age = data['age']
    gender = data['gender']
    client = MongoClient('mongodb://localhost:27017/')
    db = client['medic_app']
    new_user = {'id': username, 'pwd': password, 'age': age, 'gender': gender}
    if role == '0':
        collection = db['doctor']
        if collection.find_one({'id': username}):
            return JsonResponse({"permission": False})
        collection.insert_one(new_user)
    elif role == '1':
        collection = db['patient']
        if collection.find_one({'id': username}):
            return JsonResponse({"permission": False})
        collection.insert_one(new_user)
    return JsonResponse({"permission": True})

def generate(request, p_id):
    data = request.POST
    print(data)
    message = data['message']
    doctor_id = data['doctor_id']

    custom_settings = init_prompts()

    my_record = inference(message, custom_settings)

    client = MongoClient('mongodb://localhost:27017/')
    db = client['medic_app']
    collection = db['record']


    record = {'p_id': p_id,
              'd_id': doctor_id,
              'name': my_record['姓名'],
              'age' : my_record['年龄'],
              'chief': my_record['主诉'],
              'lastingTime': my_record['持续时间'],
              'history': my_record['既往病史'],
              'allergy': my_record['过敏药物']
              }

    collection.insert_one(record)

    return JsonResponse({"message": my_record}, json_dumps_params={'ensure_ascii': False})

def myCase(request, p_id):
    data = request.POST
    client = MongoClient('mongodb://localhost:27017/')
    db = client['medic_app']
    record = db['record']
    cases = record.find({'p_id': p_id}, {'_id': 0})
    case_list = list(cases)

    return JsonResponse({'case_list': case_list})


def diagnostic(request, d_id):
    client = MongoClient('mongodb://localhost:27017/')
    db = client['medic_app']
    record = db['record']
    cases = record.find({'d_id': d_id}, {'_id': 0})
    case_list = list(cases)
    return JsonResponse({'patientList': case_list})


def case(request, p_id, d_id):
    client = MongoClient('mongodb://localhost:27017/')
    db = client['medic_app']
    record = db['record']
    cases = record.find({'p_id': p_id, 'd_id': d_id}, {'_id': 0})
    case_list = list(cases)
    return JsonResponse({'case_list': case_list})


schema = {
    '病历条目': ['姓名', '年龄', '主诉', '持续时间', '既往病史', '过敏药物']
}

IE_PATTERN = "{}\n\n提取上述句子中{}类型的实体，输出JSON格式，如果上述句子没有该信息，则用['无']来表示，多个值之间用','分隔。"

ie_examples = {
    '病历1':
        {
            'sentence': '我叫傅一帆，今年38岁，我从昨天开始感觉头疼，并且伴有腹泻的症状。以前得过肠胃型感冒，对阿莫西林过敏',
            'answers': {
                '姓名': '傅一帆',
                '年龄': '38岁',
                '主诉': '头疼，腹泻',
                '持续时间': '1天',
                '既往病史': '肠胃型感冒',
                '过敏药物': '阿莫西林',
            }
        },
    '病历2':
        {
            'sentence': '我叫林凡，今年42岁，前天开始一直觉得肩膀酸痛，几十年前曾经得过肩周炎，但是已经痊愈，没有发现药物过敏',
            'answers': {
                '姓名': '林凡',
                '年龄': '42岁',
                '主诉': '肩膀酸痛',
                '持续时间': '2天',
                '既往病史': '肩周炎',
                '过敏药物': '未提供',
            }
        },
    '病历3':
        {
            'sentence': '我叫何清，今年13岁，这半个月一直觉得脚疼',
            'answers': {
                '姓名': '何清',
                '年龄': '13岁',
                '主诉': '脚疼',
                '持续时间': '半个月',
                '既往病史': '未提供',
                '过敏药物': '未提供',
            }
        },
    '病历4':
        {
            'sentence': '我叫林露，我觉得牙疼，去年因为口腔癌做过手术',
            'answers': {
                '姓名': '林露',
                '年龄': '未提供',
                '主诉': '牙疼',
                '持续时间': '未提供',
                '既往病史': '口腔癌',
                '过敏药物': '未提供',
            }
        }
}


def init_prompts():
    ie_prefix = [
        (
            "需要你协助完成病历信息抽取任务，当我给你一个病人自述时，帮我抽取出句子中的‘姓名、年龄、主诉、持续时间、既往病史、过敏药物’，并按照\
            JSON的格式输出，如果缺少信息用['未提供']来表示，多个值之间用','分隔。持续时间要说明具体几天或几个月，只要输出json格式数据，不要输出\
            多余信息",
            '好的，我将只以json格式输出。'
        )
    ]
    properties_str = ', '.join(schema['病历条目'])

    schema_str_list = f'“病历条目”({properties_str})'

    for key, example in ie_examples.items():
        sentence = example['sentence']
        formatted_prompt = IE_PATTERN.format(sentence, schema_str_list)
        answers_json = json.dumps(example['answers'], ensure_ascii=False)
        ie_prefix.append((formatted_prompt, answers_json))

    return {'ie_prefix': ie_prefix}


def format_output(response: str):
    if '```json' in response:
        # response = response.replace("\n","")
        res = re.findall(r'```json((.|\n)*?)```', response)
        if len(res) and res[0]:
            response = res[0]
            response = response[0]
        response.replace('、', ',')
    elif '{' in response:
        res = re.findall(r'{((.|\n)*?)}', response)
        if len(res) and res[0]:
            response = res[0]
            response = response[0]
            response = '{' + response + '}'
        response.replace('、', ',')

    try:
        return json.loads(response)
    except:
        return response


def inference(sentence, custom_settings: dict):
    properties_str = ', '.join(schema['病历条目'])
    schema_str_list = f'“病历条目”({properties_str})'
    sentence_with_ie_prompt = IE_PATTERN.format(sentence, schema_str_list)
    result, _ = (model.chat(tokenizer, sentence_with_ie_prompt, history=custom_settings['ie_prefix']))
    result = format_output(result)
    return result
